#include <stdio.h>
#include <string.h>
#include <math.h>
#include "common.h"

#define MAX(x, y) (((x) > (y)) ? (x) : (y))
#define MIN(x, y) (((x) < (y)) ? (x) : (y))

#define CPUTILESIZE 64


/*
 *  CUDA version for two vector arrays A and B
 */
__global__ void corrcoefmatrixAB_cuda(float *out, float *inA, float *inB, int n, int m) {
  __shared__ float inAS[TILESIZE][TILESIZE];
  __shared__ float inBS[TILESIZE][TILESIZE];

  int i,j,k;
  int tileBegin;
  float a_x,b_x;
  float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

  i = blockIdx.y*blockDim.y + threadIdx.y;
  j = blockIdx.x*blockDim.x + threadIdx.x;

  sum_a = sum_a2 = sum_b = sum_b2 = sum_ab = 0;
  for (tileBegin=0; tileBegin < m; tileBegin += blockDim.x) {
     // Load data into shared memory
     inAS[threadIdx.y][threadIdx.x] = inA[(blockIdx.y*blockDim.y + threadIdx.y)*m+tileBegin+threadIdx.x];
     inBS[threadIdx.y][threadIdx.x] = inB[(blockIdx.x*blockDim.x + threadIdx.y)*m+tileBegin+threadIdx.x];
     __syncthreads();

     // Do partial calculation on each tile
     for (k=0; k < blockDim.x; k++) {
        if (tileBegin + k < m) { // we should test if it is as good as padding
           a_x = inAS[threadIdx.y][k];
           b_x = inBS[threadIdx.x][k];
           sum_a += a_x;
           sum_a2 += a_x*a_x;
           sum_b += b_x;
           sum_b2 += b_x*b_x;
           sum_ab += a_x*b_x;
        }
     }
     __syncthreads();
  }
  corrcoef = (m*sum_ab - sum_a*sum_b)/sqrtf((m*sum_a2-sum_a*sum_a)*(m*sum_b2-sum_b*sum_b));
  out[i*n+j] = corrcoef;
}

/*
 *  OpenMP version for two vector arrays A and B
 */
void corrcoefmatrixAB_omp(float *out, float *inA, float *inB, int n, int m) {

  int i,j,k;

  #pragma omp parallel for private(j,k) 
  for (i=0; i < n; i++) {

     float a_x,b_x;
     float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

     for (j=0; j < n; j++) {
        sum_a = sum_a2 = sum_b = sum_b2 = sum_ab = 0;
        for (k=0; k < m; k++) {
           a_x = inA[i*m+k];
           b_x = inB[j*m+k]; 
           sum_a += a_x;
           sum_a2 += a_x*a_x;
           sum_b += b_x;
           sum_b2 += b_x*b_x;
           sum_ab += a_x*b_x;
        }
        corrcoef = (m*sum_ab - sum_a*sum_b)/sqrtf((m*sum_a2-sum_a*sum_a)*(m*sum_b2-sum_b*sum_b));
        out[i*n+j] = corrcoef;
     }

  }

}

/*
 *  OpenMP version
 */
void corrcoefmatrix_omp(float *out, float *in, int n, int m) {

  int i,j,k;
//  float a_x,b_x;
//  float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

  #pragma omp parallel for private(j,k) 
  for (i=0; i < n; i++) {

     float a_x,b_x;
     float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

//     #pragma omp critical
//     fprintf(stderr,"%d,%d\n",omp_get_thread_num(),i);

     for (j=0; j < n; j++) {
        sum_a = sum_a2 = sum_b = sum_b2 = sum_ab = 0;
        for (k=0; k < m; k++) {
           a_x = in[i*m+k];
           b_x = in[j*m+k]; 
           sum_a += a_x;
           sum_a2 += a_x*a_x;
           sum_b += b_x;
           sum_b2 += b_x*b_x;
           sum_ab += a_x*b_x;
        }
        corrcoef = (m*sum_ab - sum_a*sum_b)/sqrtf((m*sum_a2-sum_a*sum_a)*(m*sum_b2-sum_b*sum_b));
        out[i*n+j] = corrcoef;
     }

  }

}

/*
 *  OpenMP version with tiling for two vector arrays A and B
 */
void corrcoefmatrixAB_omp_tiled(float *out, float *inA, float *inB, int n, int m) {

  int i,j,k;
  int ii,jj;
  int last_i, last_j;

  #pragma omp parallel for private(jj,i,j,k,last_i,last_j) 
  for (ii=0; ii < n; ii += CPUTILESIZE) {
     for (jj=0; jj < n; jj += CPUTILESIZE) {
        last_i = MIN(ii+CPUTILESIZE,n);
        last_j = MIN(jj+CPUTILESIZE,n);

        for (i=ii; i < last_i; i++) {


           float a_x,b_x;
           float sum_a, sum_b, sum_a2, sum_b2, sum_ab, corrcoef;

           for (j=jj; j < last_j; j++) {
              sum_a = sum_a2 = sum_b = sum_b2 = sum_ab = 0;
              for (k=0; k < m; k++) {
                 a_x = inA[i*m+k];
                 b_x = inB[j*m+k]; 
                 sum_a += a_x;
                 sum_a2 += a_x*a_x;
                 sum_b += b_x;
                 sum_b2 += b_x*b_x;
                 sum_ab += a_x*b_x;
              }
              corrcoef = (m*sum_ab - sum_a*sum_b)/sqrtf((m*sum_a2-sum_a*sum_a)*(m*sum_b2-sum_b*sum_b));
              out[i*n+j] = corrcoef;
           }

        }

      }
   }
}

